package com.yhzdys.ocr.service.support;

import ai.djl.Device;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.repository.zoo.ZooModel;

/**
 * OCR执行线程，一个线程绑定一个Predictor
 */
public class OcrExecThread extends Thread {

    private final Predictor<Image, DetectedObjects> detector;
    private final Predictor<Image, String> recognizer;

    public OcrExecThread(Runnable runnable,
                         ZooModel<Image, DetectedObjects> detectionModel,
                         ZooModel<Image, String> recognitionModel) {
        super(runnable);
        super.setName("ocr-exec-" + super.getId());
        Predictor<Image, DetectedObjects> detector = detectionModel.newPredictor(Device.cpu());
        Predictor<Image, String> recognizer = recognitionModel.newPredictor(Device.cpu());
        this.detector = detector;
        this.recognizer = recognizer;
    }

    public Predictor<Image, DetectedObjects> getDetector() {
        return detector;
    }

    public Predictor<Image, String> getRecognizer() {
        return recognizer;
    }

    @Override
    protected void finalize() throws Throwable {
        try {
            detector.close();
            recognizer.close();
        } finally {
            super.finalize();
        }
    }
}
